import pandas as pd
from torch.utils.data import Dataset

class ContentSet(Dataset):
    def __init__(self, filename, label_type=0):
        super(ContentSet, self).__init__()
        if label_type == 0:
            self.data = pd.read_csv(filename)[['id', 'cat_id', 'content', 'ncw_label']].values.tolist()
        elif label_type ==-1: # for category classification
            self.data = pd.read_csv(filename)[['id', 'cat_id', 'content']].values.tolist()
        elif label_type ==-2: # for the inference of category
            self.data = pd.read_csv(filename)[['id', 'content']].values.tolist()
        else:
            df = pd.read_csv(filename)
            df['label'] = df[['ncw_label', 'fake_label', 'real_label']].values.argmax(axis=1)
            self.data = df[['id', 'cat_id', 'content', 'label']].values.tolist()

    def __getitem__(self, index):
        return self.data[index]
    def __len__(self):
        return len(self.data)
